import numpy as np
import mindspore as ms
import mindspore.nn as nn
from mindspore.common.initializer import initializer
from mindspore.train.serialization import load_param_into_net
from mindspore import Parameter


class mindspore_init():
    def __init__(self):
        self.with_bias_cell = [nn.Conv3d, nn.Conv1d, nn.Conv2d, nn.Dense]
        self.rnn_cell = [nn.GRU, nn.LSTM, nn.RNN]
    
    def mindspore_apply(self, net):
        for cell in net.cells():
             self.cell_init(cell)
             
    def cell_init(self, cell):
        cell_type = type(cell)
        if cell_type in self.with_bias_cell:
            weight_shape = cell.weight.shape
            bias_shape = cell.bias.shape
            weight = initializer(self.constant, weight_shape, ms.float32)
            bias_weight = initializer(0, bias_shape,ms.float32)
            cell.weight.set_data(weight)
            cell.bias.set_data(bias_weight)
        elif cell_type == nn.SequentialCell:
            for item in cell:
                self.cell_init(item)

    # # another init fun
    # def net_init(self, net, constant):
    #     self.constant = constant
    #     self.mindspore_apply(net)
    #     return net
            
    def net_init(self, net, constant):
        param_dic = {}
        for item in net.parameters_and_names():
            k = item[0]
            shape = item[1].shape
            if "bias" in k:
                v = Parameter(initializer(0, shape, ms.float32))
                param_dic[k] = v 
            else:
                v = Parameter(initializer(constant, shape, ms.float32))
                param_dic[k] = v
        load_param_into_net(net, param_dic)
        return net
